from code_generator import CodeGeneratorAI
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.text import tokenizer_from_json
import json

# 加载已训练的模型
ai = CodeGeneratorAI()
ai.model = load_model('code_model.keras')

# 加载分词器配置
with open('tokenizer_config.json', 'r') as f:
    ai.tokenizer = tokenizer_from_json(f.read())

# 定义生成代码的函数
def generate_code(prompt, temperature=0.7):
    code = ai.generate_code(prompt, temperature=temperature)
    return code

# 示例：生成一个Python函数
prompt = "def calculate_average(numbers):"  # 代码起始提示
generated_code = generate_code(prompt, temperature=0.8)

print("生成的代码：")
print(generated_code)